﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Reflection;

namespace AssemblyResolvers
{
    /// <summary>
    /// Looks for unresolved assemblies in platform specific folders.
    /// For example, if current platform is x64 and assembly was not resolved
    /// by fusion engine, this resolver will try to find x64 specific assemblies
    /// in "x64" folder and its subfolders.
    /// </summary>
    public class ProcessArchitectureBasedResolver
    {
        static Dictionary<string, string> Arch2FolderMap;

        static ProcessArchitectureBasedResolver()
        {
            Arch2FolderMap = new Dictionary<string, string>();
            Arch2FolderMap.Add("x86".ToUpperInvariant(), "x86");
            Arch2FolderMap.Add("AMD64".ToUpperInvariant(), "x64");
        }


        IList<string> FindExecutables(string name, string rootDir)
        {
            rootDir = Path.GetFullPath(rootDir);

            //specified directory does not exist
            List<string> files = new List<String>();
            if (!Directory.Exists(rootDir)) return files;

            //get executables in current directory
            string[] dlls = Directory.GetFiles(rootDir, name + ".dll");
            string[] exes = Directory.GetFiles(rootDir, name + ".exe");

            foreach (string dll in dlls) files.Add(Path.Combine(rootDir, dll));
            foreach (string exe in exes) files.Add(Path.Combine(rootDir, exe));

            foreach (string dir in Directory.GetDirectories(rootDir))
            {
                string dirFullPath = Path.Combine(rootDir, dir);

                files.AddRange(FindExecutables(name, dirFullPath));
            }

            return files;
        }

        static string GetArchictectureDir()
        {
            string arch = Environment.GetEnvironmentVariable("PROCESSOR_ARCHITECTURE");
            if (arch == null)
            {
                Trace.TraceWarning("PROCESSOR_ARCHITECTURE environment variable is not defined");
                return null;
            }

            arch = arch.ToUpperInvariant();

            if (!Arch2FolderMap.Keys.Contains(arch))
            {
                Trace.TraceWarning("PROCESSOR_ARCHITECTURE == {0}. Architecture is not recognized.", arch);
                return null;
            }

            return Arch2FolderMap[arch];
        }

        public Assembly Resolve(string assemblyFullName)
        {
            string rootArchDir = GetArchictectureDir();
            if (rootArchDir == null)
            {
                Trace.TraceWarning("Failed to determine architecture folder.");
                return null;
            }

            string[] parts = assemblyFullName.Split(new char[] { ',' });
            if (parts.Length != 4) return null;

            string shortName = parts[0].Trim();
            if (shortName == String.Empty) return null;

            IList<string> files = FindExecutables(shortName, rootArchDir);

            foreach (string file in files)
            {
                Assembly assembly = Assembly.LoadFile(Path.GetFullPath(file));
                if (assembly.FullName == assemblyFullName)
                {
                    Trace.TraceInformation("Assembly \"{0}\" is loaded from \"{1}\"", assemblyFullName, file);
                    return assembly;
                }
            }

            return null;
        }

        static public void Activate()
        {
            ProcessArchitectureBasedResolver resolver = new ProcessArchitectureBasedResolver();
            AppDomain.CurrentDomain.AssemblyResolve += (a, b) => { return resolver.Resolve(b.Name); };
        }
    }
}
